#include <sys/types.h>

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <sys/time.h>
#include <sys/param.h>
#include <sys/stat.h>
#include <sys/queue.h>
#include <sys/tree.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <err.h>
#include <unistd.h>
#include <assert.h>
#include <time.h>

#include <event.h>
#include <evdns.h>
#include <evhttp.h>

#include "amproxy.h"
#include "utils.h"
#include "status.h"
#include "log.h"

/* tell our callers that the name could not be resolved */
static void dns_dispatch_error(struct dns_cache *);
static void dns_dispatch_requests(struct dns_cache *dns_entry);
static void inform_domain_notfound(struct evhttp_request *request);
/* static void inform_proxy_settings(struct evhttp_request *request); */
static void inform_no_referer(struct evhttp_request *request);

int debug = 1;

/* globals */

static int log_fd = -1;
static short http_port = 8080;	/* port on which the server runs */

static int
dns_compare(struct dns_cache *a, struct dns_cache *b)
{
	return strcasecmp(a->name, b->name);
}

//dns缓存分裂树
static SPLAY_HEAD(dns_tree, dns_cache) root;

SPLAY_PROTOTYPE(dns_tree, dns_cache, node, dns_compare);
SPLAY_GENERATE(dns_tree, dns_cache, node, dns_compare);

static void
dns_ttl_expired(int result, short what, void *arg)
{
	struct dns_cache *dns = arg;
	
	fprintf(stderr, "Expiring dns entry for %s\n", dns->name);

	assert(TAILQ_FIRST(&dns->entries) == NULL);
	dns_free(dns);
}

static void
dns_resolv_cb(int result, char type, int count, int ttl,
    void *addresses, void *arg)
{
	struct dns_cache *entry = arg;
	struct timeval tv;

	if (result != DNS_ERR_NONE) {
		/* we were not able to resolve the name */
		dns_dispatch_error(entry);
		return;
	}


	/* copy the addresses */
	entry->addresses = calloc(count, sizeof(struct in_addr));
	if (entry->addresses == NULL)
		err(1, "calloc");
	entry->address_count = count;
	memcpy(entry->addresses, addresses, count * sizeof(struct in_addr));

	dns_dispatch_requests(entry);

	/* expire it after its time-to-live is over */
	evtimer_set(&entry->ev_timeout, dns_ttl_expired, entry);
	timerclear(&tv);
	tv.tv_sec = ttl;
	evtimer_add(&entry->ev_timeout, &tv);
}

static void
http_copy_headers(struct evkeyvalq *dst, struct evkeyvalq *src)
{
	struct evkeyval *kv;
	TAILQ_FOREACH(kv, src, next) {
		if (strcasecmp(kv->key, "Transfer-Encoding") == 0 ||
		    strcasecmp(kv->key, "Connection") == 0 ||
		    strcasecmp(kv->key, "Keep-Alive") == 0 ||
		    strcasecmp(kv->key, "Proxy-Connection") == 0) {
			if (debug)
				fprintf(stderr, "Ignoring %s: %s\n",
				    kv->key, kv->value);
			continue;
		}
		/* we might want to do some filtering here */
		if (debug)
			fprintf(stderr, "Copying %s: %s\n",
			    kv->key, kv->value);
		evhttp_add_header(dst, kv->key, kv->value);
	}
}

static void
map_location_header(struct evhttp_request *req, const char *location)
{
	static char path[1024];
	char *host, *uri;
	u_short port;
			
	if (location[0] == '/') {
		if (http_hostportfile(req->uri, &host, &port, &uri) == -1)
			return;
		snprintf(path, sizeof(path),
		    "http://%s%s", host, location);
	} else {
		strlcpy(path, location, sizeof(path));
	}
	fprintf(stderr, "Mapping %s -> %s\n", path, req->uri);
	site_new(path, req->uri);
}

static void
inform_dangerous(struct evhttp_request *request)
{
	struct evbuffer *databuf = evbuffer_new();
	char *escaped = evhttp_htmlescape(request->uri);
	assert(databuf != NULL);

	evbuffer_add_printf(databuf,
	    "<html><head><title>Dangerous Request</title></head>"
	    "<body>Possibly dangerous load of %s</body></html>",
	    escaped);
	free(escaped);

	/* we cannot allow this request */
	evhttp_send_reply(request,
	    HTTP_NOTFOUND, "Disallowing dangerous request.",
	    databuf);
	evbuffer_free(databuf);
}

static void
inform_domain_notfound(struct evhttp_request *request)
{
	struct evbuffer *databuf = evbuffer_new();
	char *escaped = evhttp_htmlescape(request->uri);
	assert(databuf != NULL);

	evbuffer_add_printf(databuf,
	    "<html><head><title>Domain not found</title></head>"
	    "<body><h1>Domain not found</h1>\n"
	    "Cannot find an IP address for %s</body></html>",
	    escaped);
	free(escaped);

	/* we cannot allow this request */
	evhttp_send_reply(request,
	    HTTP_BADREQUEST, "Disallowing dangerous request.",
	    databuf);
	evbuffer_free(databuf);
}

/*
static void
inform_proxy_settings(struct evhttp_request *request)
{
	struct evbuffer *databuf = evbuffer_new();
	assert(databuf != NULL);

	evbuffer_add_printf(databuf,
	    "<html><head><title>Proxy Configuration</title></head>"
	    "<body><div style=\"border: solid 1px; padding: 2px; "
	    "width: 40%%; "
	    "background-color: #dcdcee; font-family: Verdana, Arial;\">"
	    "<h2>Proxy Configuration</h2>\n"
	    "To use <i>amproxy</i>, you need to configure your browser to "
	    "use the following proxy <b>http://%s/</b></div><p>"
	    "<iframe src=\"http://www.amproxy.org/\" width=100%% height=85%%>"
	    "</body></html>",
	    evhttp_find_header(request->input_headers, "Host"));

	evhttp_send_reply(request, HTTP_OK, "OK", databuf);
	evbuffer_free(databuf);
}
*/
static void
inform_no_referer(struct evhttp_request *request)
{
	struct evbuffer *databuf = evbuffer_new();
	char *escaped = evhttp_encode_uri(request->uri);
	char *html_escaped = evhttp_htmlescape(request->uri);
	assert(databuf != NULL);

	evbuffer_add_printf(databuf,
	    "<html><head><title>Request Denied</title></head>"
	    "<body><div style=\"border: solid 1px; padding: 2px; "
	    "width: 100%%; "
	    "background-color: #dcdcee; font-family: Verdana, Arial;\">"
	    "<h2>Request Denied</h2>\n"
	    "To use <i>amproxy</i>, visit "
	    "<a href=\"http://amproxy.googlecode.com/\">http://amproxy.googlecode.com/</a> "
	    "or go to <a href=\"http://amproxy/?url=%s\">%s</a> "
	    "using amproxy.</div><p>"
	    "</body></html>",
	    escaped, html_escaped);
	free(escaped);
	free(html_escaped);

	/* we cannot allow this request */
	evhttp_send_reply(request, HTTP_NOTFOUND, "Not Found", databuf);
	evbuffer_free(databuf);
}

static void
http_request_done(struct evhttp_request *req, void *arg)
{
	struct site *site = NULL;
	struct proxy_request *pr = arg;
	const char *location = NULL;

	if (req == NULL || req->response_code == 0 ||
	    (site = site_find(pr->req->uri)) == NULL) {
		evhttp_send_error(pr->req,
		    HTTP_SERVUNAVAIL, "Unknown error");
		proxy_request_free(pr);
		return;
	}

	/* store the data that we received */
	site_cache_data(site, req);


	/* get the site state for filter for danger */
	if (site->danger == DANGEROUS) {
		inform_dangerous(pr->req);
		return;
	}

	location = evhttp_find_header(req->input_headers, "Location");
	/* keep track of the redirect so that we can tie it together */
	if (location != NULL)
		map_location_header(pr->req, location);
	http_copy_headers(pr->req->output_headers, req->input_headers);

	/* make everything we do no-cacheable */
	evhttp_remove_header(pr->req->output_headers, "Pragma");
	evhttp_add_header(pr->req->output_headers,
	    "Pragma", "no-cache, no-store");

	evhttp_remove_header(pr->req->output_headers, "Cache-Control");
	evhttp_add_header(pr->req->output_headers,
	    "Cache-Control", "no-cache, no-store, must-revalidate, max-age=-1");

	evhttp_send_reply(pr->req,
	    req->response_code, req->response_code_line,
	    req->input_buffer);
	proxy_request_free(pr);
}

static void
dispatch_single_request(struct dns_cache *dns, struct proxy_request *pr)
{
	struct evhttp_request *request;
	char *address = inet_ntoa(dns->addresses[0]);

	assert(pr->evcon == NULL);
	pr->evcon = evhttp_connection_new(address, pr->port);
	fprintf(stderr, "Making connection to %s:%d\n", address, pr->port);
	if (pr->evcon == NULL)
		goto fail;

	/* we got the connection now - queue the request */
	request = evhttp_request_new(http_request_done, pr);
	if (request == NULL)
		goto fail;

	http_copy_headers(request->output_headers, pr->req->input_headers);
	evhttp_add_header(request->output_headers,
	    "X-Forwarded-For", pr->req->remote_host);

	evhttp_add_header(request->output_headers, "Connection", "close");
	evhttp_make_request(pr->evcon, request, EVHTTP_REQ_GET, pr->uri);
	return;

fail:
	evhttp_send_error(pr->req,
	    HTTP_SERVUNAVAIL, "Out of resources");
	proxy_request_free(pr);
	return;
}

static void
dns_dispatch_requests(struct dns_cache *dns)
{
	struct proxy_request *entry;
	while ((entry = TAILQ_FIRST(&dns->entries)) != NULL) {
		TAILQ_REMOVE(&dns->entries, entry, next);
		
		dispatch_single_request(dns, entry);
	}
}

static void
dns_dispatch_error(struct dns_cache *dns_entry)
{
	struct proxy_request *entry;
	while ((entry = TAILQ_FIRST(&dns_entry->entries)) != NULL) {
		TAILQ_REMOVE(&dns_entry->entries, entry, next);

		inform_domain_notfound(entry->req);
		proxy_request_free(entry);
	}

	/* no negative caching */
	dns_free(dns_entry);
}

struct dns_cache *
dns_new(const char *name)
{
	struct dns_cache *entry, tmp;
	struct in_addr address;


	tmp.name = (char *)name;
	if ((entry = SPLAY_FIND(dns_tree, &root, &tmp)) != NULL)
		return (entry);

	entry = calloc(1, sizeof(struct dns_cache));
	if (entry == NULL)
		err(1, "calloc");

	entry->name = strdup(name);
	if (entry->name == NULL)
		err(1, "strdup");

	TAILQ_INIT(&entry->entries);
	SPLAY_INSERT(dns_tree, &root, entry);

	if (inet_aton(entry->name, &address) != 1) {
		evdns_resolve_ipv4(entry->name, 0,
		    dns_resolv_cb, entry);
	} else {

		/* we already have an address - no dns necessary */
		dns_resolv_cb(DNS_ERR_NONE, DNS_IPv4_A,
		    1, 3600, &address, entry);
	}

	return (entry);
}

void
dns_free(struct dns_cache *entry)
{
	SPLAY_REMOVE(dns_tree, &root, entry);
	free(entry->name);
	free(entry);
}

void
request_add_dns(struct dns_cache *entry, struct proxy_request *pr)
{
	TAILQ_INSERT_TAIL(&entry->entries, pr, next);

	/* still waiting for resolution */
	if (entry->address_count == 0)
		return;

	dns_dispatch_requests(entry);
}

/* keep some simple stats */
//简单的记录一些浏览的状态
static void
amproxy_traffic_stats(struct site *site)
{
	extern struct stats statistics;

	statistics.num_requests++;
	switch (site->danger) {
	case HARMLESS:
		statistics.num_harmless++;
		break;
	case DANGEROUS:
		statistics.num_dangerous++;
		break;
	case UNKNOWN:
	default:
		statistics.num_unknown++;
		break;
	}
}

/*
 * Receive all possible requests - analyze them for doing stuff
 */

void
request_handler(struct evhttp_request *request, void *arg)
{
	char *host, *uri;
	const char *referer;
	u_short port;
	struct site *site;
	struct dns_cache *entry;
	struct proxy_request *pr;
	

	if (http_hostportfile(request->uri, &host, &port, &uri) == -1) {
		evhttp_send_error(request,
		    HTTP_BADREQUEST, "Illegal request.");
		return;
	}

	//使用amproxy处理
	if (strncasecmp(host, "amproxy", 6) == 0) {
		fprintf(stderr, "Received request for %s\n", request->uri);
		amproxy_handle_request(request, arg);
		return;
	}

	/* now insert the request into our status object */
	//referer指向网页的引用地址即根站点
	referer = evhttp_find_header(request->input_headers, "Referer");
	fprintf(stderr, "Received request for %s (%s) from %s\n",
	    request->uri, referer, request->remote_host);
	if (referer == NULL) {
		inform_no_referer(request);
		return;
	}
	//发出浏览请求	
	site = site_new(request->uri, referer);
		
	//	if (site == NULL) {
		/* some evil sites are tricky */
	//	char *decode = evhttp_decode_uri(referer);
	//	site = site_new(request->uri, decode);
	//	free(decode);
	//}
	//找不到站点
	if (site == NULL) {
		fprintf(stderr, "Missing site state for %s\n", request->uri);
		inform_no_referer(request);
		return;
	}
	//记录状态
	amproxy_traffic_stats(site);

	/*
	 * make requests to dangerous sites just to keep track of the content.
	 * xxx: we might want to wait with logging this until we got the data.
	 */
	if (site->danger == DANGEROUS) {
		if (log_fd != -1)
			log_dangerous_request(log_fd, request);
	}

	/* make sure that we do not send a referer if this is a root URL */
	if (site->parent == NULL)
		evhttp_remove_header(request->input_headers, "Referer");

	entry = dns_new(host);
	pr = proxy_request_new(request, port, uri);
	request_add_dns(entry, pr);
}

struct proxy_request *
proxy_request_new(struct evhttp_request *req, u_short port, char *uri)
{
	struct proxy_request *pr;

	if ((pr = calloc(1, sizeof(struct proxy_request))) == NULL)
		err(1, "calloc");

	pr->uri = strdup(uri);
	if (pr->uri == NULL)
		err(1, "strdup");

	pr->req = req;
	pr->port = port;

	return (pr);
}

void
proxy_request_free(struct proxy_request *pr)
{
	/* if (pr->evcon != NULL)
	   evhttp_connection_free(pr->evcon); */

	free(pr->uri);
	free(pr);
}

void
usage(const char *progname)
{
	fprintf(stderr,
	    "%s: [-p port] [-g good] [-b bad]\n"
	    "\t -g good_patterns a file or url containing the good patterns\n"
	    "\t -b bad_patterns a file or url containing the danger patterns\n"
	    "\t for documentation of all options consult the man page\n",
	    progname);
}

//主函数
int
main(int argc, char **argv)
{
	extern char *optarg;
	extern int optind;
	struct evhttp *http_server = NULL;
	char *whitelist = "white_list";
	char *blacklist = "black_list";
	char *log_file = "amproxy.log";
	int ch;
//参数选择
	while ((ch = getopt(argc, argv, "Pp:g:b:l:R:")) != -1)
		switch(ch) {
		case 'p':
			http_port = atoi(optarg);
			if (!http_port) {
				usage(argv[0]);
				exit(1);
			}
			break;
		case 'l':
			if (!strlen(optarg))
				log_file = NULL;
			else
				log_file = optarg;
			break;
		case 'g':
			whitelist = optarg;
			break;
		case 'b':
			blacklist = optarg;
			break;
		default:
			usage(argv[0]);
			exit(1);
		}

	argc -= optind;
	argv += optind;
	//初始化
	event_init();
	evtag_init();
	evdns_init();

	fprintf(stderr, "amproxy %s starting up ...\n", VERSION);
	
	//初始化黑名单和白名单
	status_init(whitelist, blacklist);

	//启动服务器
	http_server = evhttp_start("0.0.0.0", http_port);
	if (http_server == NULL) {
		fprintf(stderr, "Cannot run web server on port %d\n",
		    http_port);
		exit(1);
	}
	evhttp_set_gencb(http_server, request_handler, NULL);

	fprintf(stderr, "Starting web server on port %d\n"
	    "Configure your browser to use this server as proxy.\n",
	    http_port);
	//初始化DNS缓存树
	SPLAY_INIT(&root);

	event_dispatch();

	if (log_fd != -1)
		log_close(log_fd);

	exit (0);
}
